#include <math.h>
#include <float.h>
#include <curand_kernel.h>

// 定义 FLT_MAX 替代标准库中的定义
#ifndef FLT_MAX
#define FLT_MAX 3.402823466e+38F
#endif

// 将 my_expf 定义放在全局作用域
__device__ inline float my_expf(float x) {
    return __expf(x);
}

__device__ float calculate_swap_delta_opt(float* distmat, int* sol, int n, int i, int j) {
    // 使用寄存器存储常用值
    const int prev_i = (i - 1 + n) % n;
    const int next_i = (i + 1) % n;
    const int prev_j = (j - 1 + n) % n;
    const int next_j = (j + 1) % n;
    
    const int sol_prev_i = sol[prev_i];
    const int sol_i = sol[i];
    const int sol_next_i = sol[next_i];
    const int sol_prev_j = sol[prev_j];
    const int sol_j = sol[j];
    const int sol_next_j = sol[next_j];
    
    // 预计算索引
    const int idx_prev_i_i = sol_prev_i * n + sol_i;
    const int idx_i_next_i = sol_i * n + sol_next_i;
    const int idx_prev_j_j = sol_prev_j * n + sol_j;
    const int idx_j_next_j = sol_j * n + sol_next_j;
    const int idx_prev_i_j = sol_prev_i * n + sol_j;
    const int idx_j_next_i = sol_j * n + sol_next_i;
    const int idx_prev_j_i = sol_prev_j * n + sol_i;
    const int idx_i_next_j = sol_i * n + sol_next_j;
    
    float delta = 0;
    delta += distmat[idx_prev_i_j];
    delta += distmat[idx_j_next_i];
    delta += distmat[idx_prev_j_i];
    delta += distmat[idx_i_next_j];
    delta -= distmat[idx_prev_i_i];
    delta -= distmat[idx_i_next_i];
    delta -= distmat[idx_prev_j_j];
    delta -= distmat[idx_j_next_j];
    
    if (j - i == 1 || (i == 0 && j == n - 1)) {
        const int idx_i_j = sol_i * n + sol_j;
        delta += 2 * distmat[idx_i_j];
    }
    
    return delta;
}

__device__ float calculate_2opt_delta(float* distmat, int* sol, int n, int i, int j) {
    int prev_i = (i - 1 + n) % n;
    int next_j = (j + 1) % n;
    float delta = distmat[sol[prev_i] * n + sol[j]];
    delta += distmat[sol[i] * n + sol[next_j]];
    delta -= distmat[sol[prev_i] * n + sol[i]];
    delta -= distmat[sol[j] * n + sol[next_j]];
    if (i == 0 && j == n - 1) {
        delta = 0;
    }
    return delta;
}

__device__ float calculate_move_delta(float* distmat, int* sol, int n, int i, int j) {
    int next_i = (i + 1) % n;
    int prev_j = (j - 1 + n) % n;
    int next_j = (j + 1) % n;
    float delta = distmat[sol[i] * n + sol[j]];
    delta += distmat[sol[j] * n + sol[next_i]];
    delta += distmat[sol[prev_j] * n + sol[next_j]];
    delta -= distmat[sol[i] * n + sol[next_i]];
    delta -= distmat[sol[prev_j] * n + sol[j]];
    delta -= distmat[sol[j] * n + sol[next_j]];
    if (i == n - 1 && j == 0) {
        delta = 0;
    }
    return delta;
}

__device__ void apply_swap(int* sol, int i, int j) {
    int temp = sol[i];
    sol[i] = sol[j];
    sol[j] = temp;
}

__device__ void apply_2opt(int* sol, int i, int j) {
    int start = i;
    int end = j;
    while (start < end) {
        int temp = sol[start];
        sol[start] = sol[end];
        sol[end] = temp;
        start++;
        end--;
    }
}

__device__ void apply_move(int* sol, int i, int j) {
    int temp = sol[j];
    for (int k = j; k > i + 1; k--) {
        sol[k] = sol[k - 1];
    }
    sol[i + 1] = temp;
}

extern "C" __global__ void init_randoms(char* states, unsigned long seed) {
    // 将 char* 转换为 curandState*
    curandState* curand_states = reinterpret_cast<curandState*>(states);
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    curand_init(seed, idx, 0, &curand_states[idx]);
}

extern "C" __global__ void sa_kernel_opt(
    float* distmat, 
    int* solutions, 
    float* costs, 
    float* temperatures,
    int max_tnm,
    int n,
    int L,
    char* random_states_char  // 添加此参数
) {
    // 将 char* 转换为 curandState*
    curandState* random_states = reinterpret_cast<curandState*>(random_states_char);
    
    int replica_idx = blockIdx.x;
    int tid = threadIdx.x;
    
    // 使用共享内存存储解决方案
    extern __shared__ int shared_sol[];
    
    // 协作加载解决方案
    for (int i = tid; i < n; i += blockDim.x) {
        shared_sol[i] = solutions[replica_idx * n + i];
    }
    __syncthreads();
    
    // 线程私有变量
    float best_delta = FLT_MAX;
    int best_i = -1;
    int best_j = -1;
    int best_op = -1;
    
    curandState local_state = random_states[replica_idx * blockDim.x + tid];
    
    // 每个线程处理多个邻域操作
    const int ops_per_thread = max_tnm / (blockDim.x * gridDim.y) + 1;
    
    for (int k = 0; k < ops_per_thread; k++) {
        int op_type = curand(&local_state) % 3;
        int i = curand(&local_state) % n;
        int j = curand(&local_state) % n;
        
        if (i > j) {
            int temp = i;
            i = j;
            j = temp;
        }
        if (i == j) {
            j = (j + 1) % n;
            if (i == j) continue;  // 避免无效对
        }
        
        float delta = 0.0f;
        switch(op_type) {
            case 0: delta = calculate_swap_delta_opt(distmat, shared_sol, n, i, j); break;
            case 1: delta = calculate_2opt_delta(distmat, shared_sol, n, i, j); break;
            case 2: delta = calculate_move_delta(distmat, shared_sol, n, i, j); break;
        }
        
        if (delta < best_delta) {
            best_delta = delta;
            best_i = i;
            best_j = j;
            best_op = op_type;
        }
    }

    __shared__ float s_deltas[256];
    __shared__ int s_i[256], s_j[256], s_op[256];
    
    s_deltas[tid] = best_delta;
    s_i[tid] = best_i;
    s_j[tid] = best_j;
    s_op[tid] = best_op;
    __syncthreads();
    
    for (int s = blockDim.x/2; s > 0; s >>= 1) {
        if (tid < s) {
            if (s_deltas[tid + s] < s_deltas[tid]) {
                s_deltas[tid] = s_deltas[tid + s];
                s_i[tid] = s_i[tid + s];
                s_j[tid] = s_j[tid + s];
                s_op[tid] = s_op[tid + s];
            }
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        // 使用标准 expf 函数
        float acceptance_prob = expf(-s_deltas[0] / temperatures[replica_idx]);
        float rand_val = curand_uniform(&local_state);
        
        if (s_deltas[0] < 0 || rand_val < acceptance_prob) {
            if (s_op[0] == 0) {
                apply_swap(shared_sol, s_i[0], s_j[0]);
            } else if (s_op[0] == 1) {
                apply_2opt(shared_sol, s_i[0], s_j[0]);
            } else {
                apply_move(shared_sol, s_i[0], s_j[0]);
            }
            costs[replica_idx] += s_deltas[0];
        }
        
        for (int i = 0; i < n; i++) {
            solutions[replica_idx * n + i] = shared_sol[i];
        }
    }
    
    random_states[replica_idx * blockDim.x + tid] = local_state;
}